import torch
import matplotlib.pyplot as plt
import numpy as np
from gen_data import gen

# x_data = torch.Tensor([[1.0], [2.0], [3.0]])
# y_data = torch.Tensor([[2.0], [4.0], [6.0]])
x_data, y_data = gen(28.75, -750)
x_data = np.array(x_data)
y_data = np.array(y_data)
scale_x = 1 / 200.0
scale_y = 1 / 5000.0
print('scale x=' + str(scale_x))
print('scale y=' + str(scale_y))
x_data *= scale_x
y_data *= scale_y
x_data = torch.Tensor(np.array(x_data))
y_data = torch.Tensor(np.array(y_data))




class LinearModel(torch.nn.Module):
    def __init__(self):  # 构造函数
        super(LinearModel, self).__init__()
        self.linear = torch.nn.Linear(1, 1)  # 构造对象，并说明输入输出的维数，第三个参数默认为true，表示用到b

    def forward(self, x):
        y_pred = self.linear(x)  # 可调用对象，计算y=wx+b
        return y_pred


model = LinearModel()  # 实例化模型

criterion = torch.nn.MSELoss(reduction='sum')
# model.parameters()会扫描module中的所有成员，如果成员中有相应的权重，那么都会将结果加到要训练的集合参数上
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)  # lr为学习率

epoch_list = []
loss_list = []
# for epoch in np.arange(0, 100, 2):
for epoch in range(100):
    y_pred = model(x_data)
    loss = criterion(y_pred, y_data)
    print(epoch, loss.item())

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    epoch_list.append(epoch)
    loss_list.append(loss.item())
print('w=', model.linear.weight.item())
print('b=', model.linear.bias.item())

x_test = torch.Tensor([[0.4]])
y_test = model(x_data)

def drawPredict():
    plt.title('Predict')
    # 蓝色是生成的原数据
    plt.plot(x_data, y_data, color = 'blue', label = 'origin',linewidth = 1.0, linestyle= '--')      #define line color and style
    # 红色是预测的数据
    plt.plot(x_data, np.reshape(y_test.data, (50)), color = 'red', label = 'predict',linewidth = 1.0, linestyle= '--')      #define line color and style
    plt.xlabel('scaled x')
    plt.ylabel('scaled y')
    plt.legend()

def drawLoss():
    plt.figure()
    # 使用SGD训练的loss与训练epoch曲线
    plt.title('SGD')
    plt.xlabel('times')
    plt.ylabel('loss')
    plt.plot(epoch_list, loss_list)
    plt.legend()

drawPredict()
drawLoss()
# print('y_pred = ', y_test.data)
plt.show()
